package top.ochiamalu.aioj.websocket;

import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import top.ochiamalu.aioj.websocket.model.JudgeRequestMessage;
import top.ochiamalu.aioj.websocket.model.JudgeResponseMessage;

import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import static top.ochiamalu.aioj.constant.JudgeConstants.RETRY_COUNT;
import static top.ochiamalu.aioj.constant.JudgeConstants.RETRY_INTERVAL;

/**
 * 判题处理器
 *
 * @author ochiamalu
 * @date 2025/01/24
 */
@Slf4j
public class JudgeHandler extends TextWebSocketHandler {

    private static final ConcurrentHashMap<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();

    private final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * 发送判题状态
     *
     * @param submissionId 提交id
     * @param type         类型
     */
    public static void sendJudgeStatus(String submissionId, Integer type) {
        // 等待WebSocket连接建立
        int retryCount = 0;
        while (!SESSIONS.containsKey(submissionId) && retryCount < RETRY_COUNT) {
            try {
                TimeUnit.MILLISECONDS.sleep(RETRY_INTERVAL);
                retryCount++;
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                return;
            }
        }

        WebSocketSession session = SESSIONS.get(submissionId);
        if (session != null && session.isOpen()) {
            try {
                JudgeResponseMessage response = new JudgeResponseMessage();
                response.setType(type);
                session.sendMessage(new TextMessage(new ObjectMapper().writeValueAsString(response)));
            } catch (Exception e) {
                log.error("发送判题状态失败", e);
            }
        } else {
            log.warn("WebSocket session not found or closed for submissionId: {}", submissionId);
        }
    }

    /**
     * 获取WebSocket会话
     *
     * @param submissionId 提交ID
     * @return WebSocket会话
     */
    public static WebSocketSession getSession(String submissionId) {
        return SESSIONS.get(submissionId);
    }

    /**
     * 连接建立后
     *
     * @param session 阶段
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        String submissionId = Objects.requireNonNull(session.getUri()).getQuery().split("=")[1];
        SESSIONS.put(submissionId, session);

        // 发送等待评测状态
        try {
            JudgeResponseMessage response = new JudgeResponseMessage();
            response.setType(0);
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(response)));
        } catch (Exception e) {
            log.error("发送等待评测状态失败", e);
        }
    }

    /**
     * 处理消息
     *
     * @param session 阶段
     * @param message 消息
     * @throws Exception 例外
     */
    @Override
    protected void handleTextMessage(@NotNull WebSocketSession session, TextMessage message) throws Exception {
        JudgeRequestMessage request = objectMapper.readValue(message.getPayload(), JudgeRequestMessage.class);
        // 处理客户端消息
        log.info("收到消息：{}", request);
    }

    /**
     * 连接关闭后
     *
     * @param session 阶段
     * @param status  地位
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, @NotNull CloseStatus status) {
        String submissionId = Objects.requireNonNull(session.getUri()).getQuery().split("=")[1];
        SESSIONS.remove(submissionId);
    }
}
